#include "space_hash_cuda.h"
#include "cuda_block_size_set.h"
#include <cuda_runtime.h>
#include <ctime>
#include <iostream>

namespace PhysLeo
{
    
template<typename T>
void SpaceHashCuda<T>::setSpace(T resolution, glm::tvec3<T> low, glm::tvec3<T> high, int np_max)
{
    this->resolution_ = resolution;
    this->low_ = low;
    this->high_ = high;
    this->np_max_ = np_max;

    glm::tvec3<T> num_segment = (this->high_ - this->low_) / this->resolution_;
    
    this->nx_ = static_cast<int>(ceil(num_segment.x));
    this->ny_ = static_cast<int>(ceil(num_segment.y));
    this->nz_ = static_cast<int>(ceil(num_segment.z));

    this->num_ = this->nx_ * this->ny_ * this->nz_;
    this->ids_ = std::make_shared<BufferData<int>>(this->num_*this->np_max_);
    this->ids_->allocateGpu();
    this->counter_ = std::make_shared<BufferData<int>>(this->num_);
    this->counter_->allocateGpu();
    std::cout << "set space with " << this->num_ << " grids and " << this->np_max_ << " count space for each grid" << std::endl;
}

template<typename T>
__global__ static void K_SpaceHashCuda_constructHashTable(glm::tvec3<T>* pos_ptr, int* counter_ptr, int* ids_ptr, glm::tvec3<T> low, T resolution, int ny, int nz, int np_max, int size)
{
    int pid = threadIdx.x + (blockIdx.x * blockDim.x);
    if (pid >= size) return;

    glm::tvec3<int> ijk = glm::floor((pos_ptr[pid] - low) / resolution);
    int gid = ijk.x*ny*nz + ijk.y*nz + ijk.z;

    if (gid < 0) return;

    int index = atomicAdd(&(counter_ptr[gid]), 1);
    index = index < np_max - 1 ? index : np_max - 1;
    ids_ptr[gid*np_max + index] = pid;
}

template<typename T>
void SpaceHashCuda<T>::constructHashTable(std::shared_ptr<BufferData<glm::tvec3<T>>> ptr_pos)
{
    this->counter_->setZeroGpu();
    this->ids_->setZeroGpu();

    const int pos_num = ptr_pos->size();
    auto pos_ptr = ptr_pos->ptrGpu();
    auto counter_ptr = this->counter_->ptrGpu();
    auto ids_ptr = this->ids_->ptrGpu();

    K_SpaceHashCuda_constructHashTable<T> << <numBlocks(pos_num), BLOCK_SIZE >> > (pos_ptr, counter_ptr, ids_ptr, this->low_, this->resolution_, this->ny_, this->nz_, this->np_max_, pos_num);
    synchronCheck;
}

template<typename T>
__global__ static void K_SpaceHashCuda_queryNeighbor(glm::tvec3<T>* pos_ptr, NeighborList* neighbor_list_ptr, int* counter_ptr, int* ids_ptr, T max_radius, T min_radius, glm::tvec3<T> low, T resolution, int num, int ny, int nz, int np_max, int size)
{
    int pid = threadIdx.x + (blockIdx.x * blockDim.x);
    if (pid >= size) return;
    auto i = pid;

    int ids[SPACE_HASH_MAX_NEIGHBOR];
    T distance[SPACE_HASH_MAX_NEIGHBOR];

    const glm::tvec3<int> gid3 = glm::floor((pos_ptr[i] - low) / resolution);

    auto counter = 0;
    for (auto dx = -1; dx <= 1; ++dx)
        for (auto dy = -1; dy <= 1; ++dy)
            for (auto dz = -1; dz <= 1; ++dz)
            {
                const int gid = (gid3.x + dx)*ny*nz + (gid3.y + dy)*nz + (gid3.z + dz);
                if (gid < 0 || gid >= num)continue;
                if (counter_ptr[gid] > np_max)counter_ptr[gid] = np_max;
                
                for (auto j = 0; j<counter_ptr[gid]; ++j)
                {
                    int dst_index = ids_ptr[gid*np_max + j];
                    T distance_ = glm::length(pos_ptr[dst_index] - pos_ptr[i]);
                    if (distance_<max_radius&&distance_>min_radius)
                    {
                        if (counter< SPACE_HASH_MAX_NEIGHBOR)
                        {
                            ids[counter] = dst_index;
                            distance[counter] = distance_;
                            counter++;
                        }
                        else
                        {
                            auto max_id = 0;
                            T max_distance = distance[0];
                            for (auto k = 1; k<counter; ++k)
                            {
                                if (max_distance<distance[k])
                                {
                                    max_distance = distance[k];
                                    max_id = k;
                                }
                            }
                            if (distance_<max_distance)
                            {
                                ids[max_id] = dst_index;
                                distance[max_id] = distance_;
                            }
                        }
                        
                    }
                }
                
            }

    neighbor_list_ptr[i].size_ = counter;
    for (auto j = 0; j < counter; ++j)
    {
        neighbor_list_ptr[i].ids_[j] = ids[j];
    }
}

template<typename T>
void SpaceHashCuda<T>::queryNeighbor(std::shared_ptr<BufferData<glm::tvec3<T>>> ptr_pos, std::shared_ptr<BufferData<NeighborList>> ptr_neighbor_list, T max_radius, T min_radius)
{
    constructHashTable(ptr_pos);
    ptr_neighbor_list->setZeroGpu();
    synchronCheck;
    const auto start = clock();

    const int pos_num = ptr_pos->size();
    auto pos_ptr = ptr_pos->ptrGpu();
    auto counter_ptr = this->counter_->ptrGpu();
    auto ids_ptr = this->ids_->ptrGpu();
    const auto neighbor_list_ptr = ptr_neighbor_list->ptrGpu();

    K_SpaceHashCuda_queryNeighbor<T> << <numBlocks(pos_num), BLOCK_SIZE >> > (pos_ptr, neighbor_list_ptr, counter_ptr, ids_ptr, max_radius, min_radius, this->low_, this->resolution_, this->num_, this->ny_, this->nz_, this->np_max_, pos_num);
    synchronCheck;
    const auto end = clock();
    std::cout << "time cost for query neighbors cpu: " << static_cast<double>(end - start) / CLOCKS_PER_SEC << std::endl;
}

template class SpaceHashCuda<float>;
template class SpaceHashCuda<double>;

}
